import imp
from PIL import Image, ImageDraw
import math
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
import xml.etree.ElementTree as ET
import pickle
import os
import cv2
from os import getcwd
import numpy as np
from PIL import Image
import shutil
import matplotlib.pyplot as plt
import glob

import imgaug as ia
from imgaug import augmenters as iaa
#from numpy.random.mtrand import random
import random

ia.seed(1)


def read_xml_annotation(in_file):
    if not os.path.exists(in_file):
        return None

    try:
        tree = ET.parse(open(in_file, mode='rb'))
    except Exception as e:
        return None
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(float(bndbox.find('xmin').text))
        xmax = int(float(bndbox.find('xmax').text))
        ymin = int(float(bndbox.find('ymin').text))
        ymax = int(float(bndbox.find('ymax').text))
        name = object.findtext('name')
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([[xmin, ymin, xmax, ymax], name])
        # print(bndboxlist)
    #if root.find('object') is None:
    #    return None
    #bndbox = root.find('object').find('bndbox')
    return bndboxlist

def read_xml_annotation1(in_file):
    if not os.path.exists(in_file):
        return None

    try:
        tree = ET.parse(open(in_file, mode='rb'))
    except Exception as e:
        return None
    root = tree.getroot()
    bndboxlist = []
    width,height = 1, 1
    for object in root.findall('size'): 
        width = float(float(object.find('width').text))
        height = float(float(object.find('height').text))

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        x1 = float(float(bndbox.find('xmin').text)/width)
        x2 = float(float(bndbox.find('xmax').text)/width)
        y1 = float(float(bndbox.find('ymin').text)/height)
        y2 = float(float(bndbox.find('ymax').text)/height)
        name = object.findtext('name')
        # print(xmin,ymin,xmax,ymax)
        w = x2-x1
        h = y2-y1
        bndboxlist.append([[x1, y1, x2, y2], name])
        #bndboxlist.append([[x1, y1, w, h], name])
        # print(bndboxlist)
    #if root.find('object') is None:
    #    return None
    #bndbox = root.find('object').find('bndbox')
    return bndboxlist


def save_xml_annotation(outfn, imgfn, obj, size = [1, 1]):
    obj_xml = ''
    for box, name in obj:
        xmin, ymin, xmax, ymax = box
        xmin = int(xmin*size[0])
        ymin = int(ymin*size[1])
        xmax = int(xmax*size[0])
        ymax = int(ymax*size[1])
        obj_xml += f'''
        <object>
            <name>{name}</name>
            <pose>Unspecified</pose>
            <truncated>0</truncated>
            <difficult>0</difficult>
            <bndbox>
                <xmin>{xmin}</xmin>
                <ymin>{ymin}</ymin>
                <xmax>{xmax}</xmax>
                <ymax>{ymax}</ymax>
            </bndbox>
        </object>'''

    out = f'''<annotation>
	<folder></folder>
	<filename></filename>
	<path>{''}</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>1924</width>
		<height>1556</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
    {obj_xml}
    </annotation>'''

    f = open(outfn, 'w')
    f.write(out)
    f.close()
    return out

# (506.0000, 330.0000, 528.0000, 348.0000) -> (520.4747, 381.5080, 540.5596, 398.6603)


def change_xml_annotation(root, image_id, new_target):
    new_xmin = new_target[0]
    new_ymin = new_target[1]
    new_xmax = new_target[2]
    new_ymax = new_target[3]

    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    xmin = bndbox.find('xmin')
    xmin.text = str(new_xmin)
    ymin = bndbox.find('ymin')
    ymin.text = str(new_ymin)
    xmax = bndbox.find('xmax')
    xmax.text = str(new_xmax)
    ymax = bndbox.find('ymax')
    ymax.text = str(new_ymax)
    tree.write(os.path.join(root, str("%06d" % (str(id) + '.xml'))))


def change_xml_list_annotation(root, image_id, new_target, saveroot, new_name, ss=1):
    in_file = open(os.path.join(root, str(image_id)), 'rb')  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    elem = tree.find('filename')
    # elem.text = (image_id.replace('.xml', '') + str("%06d" % int(id)) + '.jpg')
    xmlroot = tree.getroot()
    index = 0

    for object in xmlroot.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        # xmin = int(bndbox.find('xmin').text)
        # xmax = int(bndbox.find('xmax').text)
        # ymin = int(bndbox.find('ymin').text)
        # ymax = int(bndbox.find('ymax').text)

        new_xmin = new_target[index][0]*ss
        new_ymin = new_target[index][1]*ss
        new_xmax = new_target[index][2]*ss
        new_ymax = new_target[index][3]*ss

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    # tree.write(os.path.join(saveroot, elem.text.replace('jpg', 'xml')), encoding='utf-8')
    tree.write(os.path.join(saveroot, new_name), encoding='utf-8')


def mkdir(path):
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        #print(path + ' 创建成功')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        #print(path + ' 目录已存在')
        return False


def is_new_then(src, out):
    if os.path.exists(out):
        out_mtime = os.path.getmtime(out)
        src_mtime = os.path.getmtime(src)
        if src_mtime > out_mtime:
            return True
    else:
        return True

    return False


def rand(a=0, b=1):
    return np.random.rand() * (b - a) + a


def sheyu(img, hue=.1, sat=1.5, val=1.5):
    # 进行色域变换
    hue = rand(-hue, hue)
    sat = rand(1, sat) if rand() < .5 else 1 / rand(1, sat)
    val = rand(1, val) if rand() < .5 else 1 / rand(1, val)
    x = rgb_to_hsv(np.array(img) / 255.)
    x[..., 0] += hue
    x[..., 0][x[..., 0] > 1] -= 1
    x[..., 0][x[..., 0] < 0] += 1
    x[..., 1] *= sat
    x[..., 2] *= val
    x[x > 1] = 1
    x[x < 0] = 0
    img = hsv_to_rgb(x)*255
    return img.astype('uint8')


def sometimes(aug): return iaa.Sometimes(0.5, aug)  # 建立lambda表达式，

import json

def dict2list(dd):
    return [(x, dd[x]) for x in dd.keys()]


def cv_imread(img_fn, flag):
    return cv2.imdecode(np.fromfile(img_fn, dtype=np.uint8), flag)

def xywhn2xyxy(bnds, size):
    outbnds = []
    for box,name in bnds:
        x,y,w,h = box
        x-=w*0.5
        y-=h*0.5
        outbnds.append([[x*size[0], y*size[1], (x+w)*size[0], (y+h)*size[0]], name])
    return outbnds


# AUGLOOP = 10  # 每张影像增强的数量
def augmentation(org_dir, out_dir, AUGLOOP=10, maxpic=10000):
    # IMG_DIR = f"{org_dir}/JPEGImages"
    # XML_DIR = f"{org_dir}/Annotations"
    outlabs = {}
    outlabs_fn = f'{out_dir}/labels.json'
    if os.path.exists(outlabs_fn):
        outlabs = json.load(open(outlabs_fn, 'r'))
        return dict2list(outlabs)

    if 0:
        for fn in outlabs.keys():
            bnds = outlabs[fn]
            xml_fn = fn.replace('.jpg', '.xml')
            img = cv_imread(fn, 1)
            h, w = img.shape[:2]
            bnds1 = xywhn2xyxy(bnds, [w, h])
            save_xml_annotation(xml_fn, fn, bnds1)

    img2jpg(org_dir)
    IMG_DIR = org_dir
    XML_DIR = org_dir

    AUG_XML_DIR = out_dir
    AUG_IMG_DIR = out_dir

    mkdir(out_dir)
    #mkdir(AUG_IMG_DIR)
    # path_file_number = len(glob.glob(pathname=IMG_DIR + '/*.jpeg'))
    #print("XML_DIR", XML_DIR)
    #path_file_number = len(glob.glob(pathname=XML_DIR + '/*.xml'))
    # print(path_file_number)
    #print("path_file_number", path_file_number)

    rotate_a = 20
    do_aug = True
    
    lab_fn_list = os.listdir(XML_DIR)
    lab_fn_list = [x for x in lab_fn_list if '.xml' in x]
    if len(lab_fn_list)>maxpic:
        random.shuffle(lab_fn_list)
        lab_fn_list = lab_fn_list[:maxpic]
        do_aug = False

    if 'ocr' in org_dir:
        rotate_a = 20
    iaa_str = 'blur flip'
    aa = []
    if 'flip' in iaa_str:
        aa = [
            iaa.Fliplr(0.5),  # 对50%的图像进行镜像翻转
            iaa.Flipud(0.5)
        ]  # 对20%的图像做左右翻转

    aa += [
        iaa.Multiply((0.8, 1.3)),  # change brightness, doesn't affect BBs
        iaa.GaussianBlur(sigma=(0, 0.5)),  # iaa.GaussianBlur(0.5),
        # 这里沿袭我们上面提到的sometimes，对随机的一部分图像做crop操作
        # crop的幅度为0到10%
        # 每个像素随机加减-10到10之间的数
        iaa.Add((-10, 10)),
        # 像素乘上0.5或者1.5之间的数字.
        #iaa.Multiply((0.5, 1.5), per_channel=0.5),
        # 对比度
        iaa.ContrastNormalization((0.8, 1.3)),
        # 将RGB变成灰度图然后乘alpha加在原图上
        iaa.Grayscale(alpha=(0.0, 1.0)),
        # 锐化处理
        iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 1.5))
    ]

    if 'Crop' in iaa_str:
        aa.append(iaa.Crop(percent=(0, 0.1)))

    if 'Affine' in iaa_str:
        aa.append(iaa.Affine(  # 对一部分图像做仿射变换
                scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},  # 图像缩放为80%到120%之间
                translate_percent={
                    "x": (-0.2, 0.2), "y": (-0.2, 0.2)},  # 平移±20%之间
                rotate=(-rotate_a, rotate_a),  # 旋转±45度之间
                shear=(-16, 16),  # 剪切变换±16度，（矩形变平行四边形）
                order=[0, 1],  # 使用最邻近差值或者双线性差值
                cval=(0, 255)  # 全白全黑填充
                # mode=ia.ALL  # 定义填充图像外区域的方法
            ))

    if 'Dropout' in iaa_str:
        aa.append(
            # 将1%到10%的像素设置为黑色
            # 或者将3%到15%的像素用原图大小2%到5%的黑色方块覆盖
            iaa.OneOf([
                iaa.Dropout((0.01, 0.1)),
                iaa.CoarseDropout(
                    (0.03, 0.15), size_percent=(0.02, 0.05)
                ),
            ]))

    if 'blur' in iaa_str:
        aa.append(
            # 将部分图像进行超像素的表示。o(╥﹏╥)o用超像素增强作者还是第一次见，比较孤陋寡闻
            #sometimes(iaa.Superpixels(p_replace=(0, 1.0), n_segments=(20, 200))),
            # 用高斯模糊，均值模糊，中值模糊中的一种增强。注意OneOf的用法
            iaa.OneOf([
                iaa.GaussianBlur((0, 3.0)),
                # 核大小2~7之间，k=((5, 7), (1, 3))时，核高度5~7，宽度1~3
                iaa.AverageBlur(k=(2, 7)),
                iaa.MedianBlur(k=(3, 11)),
            ]))

    if 'Emboss' in iaa_str:
        aa.append(
            # 浮雕效果
            iaa.Emboss(alpha=(0, 1.0), strength=(0, 2.0)))

    if 'Edge' in iaa_str:
        aa.append(
            # 边缘检测，将检测到的赋值0或者255然后叠在原图上
            iaa.OneOf([
                iaa.EdgeDetect(alpha=(0, 0.7)),
                iaa.DirectedEdgeDetect(
                    alpha=(0, 0.7), direction=(0.0, 1.0)
                ),
            ]))

    if 'Noise' in iaa_str:
        aa.append(
            # 加入高斯噪声
            iaa.AdditiveGaussianNoise(
                loc=0, scale=(0.0, 0.05*255)
            ))

    if 'Dropout' in iaa_str:
        aa.append(
            # 将1%到10%的像素设置为黑色
            # 或者将3%到15%的像素用原图大小2%到5%的黑色方块覆盖
            iaa.OneOf([
                iaa.Dropout((0.01, 0.1)),
                iaa.CoarseDropout(
                    (0.03, 0.15), size_percent=(0.02, 0.05)
                ),
            ]))

    if 'Invert' in iaa_str:
        aa.append(
            # 5%的概率反转像素的强度，即原来的强度为v那么现在的就是255-v
            #iaa.Invert(0.5, per_channel=True)
            iaa.Invert(0.5)
            )

    if 'Elastic' in iaa_str:
        aa.append(
            # 把像素移动到周围的地方。这个方法在mnist数据集增强中有见到
                iaa.ElasticTransformation(alpha=(0.5, 3.5), sigma=0.25)
            )
    if 'PiecewiseAffine' in iaa_str:
        aa.append(
            # 使用下面的0个到5个之间的方法去增强图像。注意SomeOf的用法
            # 扭曲图像的局部区域
            iaa.PiecewiseAffine(scale=(0.01, 0.05))
        )

    # 影像增强
    aa = [sometimes(x) for x in aa]
    seq = iaa.Sequential(aa)

    from tqdm import tqdm
    pbar = tqdm(lab_fn_list, total=len(lab_fn_list))

    index = 0
    for fn in pbar:
        index += 1
        fn1,ext = os.path.splitext(fn)
        if ext not in ['.xml', '.XML']:
            continue

        xml_fn = f'{XML_DIR}/{fn}'
        img_fn = f'{XML_DIR}/{fn1}.jpg'
        if not os.path.exists(xml_fn) or not os.path.exists(img_fn):
            continue

        bndbox = None
        img0 = None
        pbar.set_description(fn1)

        for epoch in range(AUGLOOP):
            out_img_fn = f'{AUG_IMG_DIR}/{fn1}_{epoch}.jpg'

            if out_img_fn in outlabs.keys() and is_new_then(out_img_fn, img_fn):
                continue
                
            if bndbox is None:
                bndbox = read_xml_annotation(xml_fn)

            if bndbox is None or len(bndbox) == 0:
                continue
            
            if img0 is None:
                img0 = cv2.imdecode(np.fromfile(img_fn, dtype=np.uint8), cv2.IMREAD_COLOR)

            if img0 is None or len(img0) == 0:
                continue
            tt = 0
            if do_aug and epoch >= tt:
                seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            # 读取图片
            # sp = img.size
            img = img0

            if random.random() > 0.5:
                #img = 255-img
                pass

            if 'sheyu' in iaa_str:
                if random.random() > 0.1:
                    img = sheyu(img)

            new_bndbox_list = []
            # bndbox 坐标增强
            if do_aug and epoch >= tt:
                bbs = None
                for i in range(len(bndbox)):
                    bnd,text = bndbox[i]
                    bbs = ia.BoundingBoxesOnImage([
                        ia.BoundingBox(
                            x1=bnd[0], y1=bnd[1], x2=bnd[2], y2=bnd[3]),
                    ], shape=img.shape)

                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]

                    # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                    n_x1 = int(
                        max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                    n_y1 = int(
                        max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                    n_x2 = int(
                        max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                    n_y2 = int(
                        max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                    if n_x1 == 1 and n_x1 == n_x2:
                        n_x2 += 1
                    if n_y1 == 1 and n_y2 == n_y1:
                        n_y2 += 1
                    if n_x1 >= n_x2 or n_y1 >= n_y2:
                        print('error', fn)
                    new_bndbox_list.append(([n_x1, n_y1, n_x2, n_y2], text))

                image_aug = seq_det.augment_images([img])[0]
                if bbs is not bbs:
                    image_auged = bbs.draw_on_image(image_aug, thickness=0)
                else:
                    image_auged = image_aug
            else:
                new_bndbox_list = bndbox
                image_auged = img
            # 存储变化后的图片
            tt = max(image_auged.shape[1], image_auged.shape[0])
            if tt>1000:
                ss = 640/max(image_auged.shape[0], image_auged.shape[1])
                ss = min(ss, 1)
                newsz = (
                    int(ss*image_auged.shape[1]), int(ss*image_auged.shape[0]))
                image_auged = cv2.resize(image_auged, newsz)
            else:
                ss = 1
            
            for i in range(len(new_bndbox_list)):
                bnd,name = new_bndbox_list[i]
                for j in range(len(bnd)):
                    bnd[j]*=ss
                    
            out_xml_fn = f'{AUG_IMG_DIR}/{fn1}_{epoch}.xml'
            save_xml_annotation(out_xml_fn, out_img_fn, new_bndbox_list)
            if True:
                height, width = image_auged.shape[:2]
                for i in range(len(new_bndbox_list)):
                    bnd,name = new_bndbox_list[i]
                    xmin = float(bnd[0]) / width
                    ymin = float(bnd[1]) / height
                    xmax = float(bnd[2]) / width
                    ymax = float(bnd[3]) / height
                    x_center = (xmin + xmax)/2
                    y_center = (ymin + ymax)/2
                    ww = xmax - xmin
                    hh = ymax - ymin
                    #if name not in classes:
                    #    classes.append(name)
                    #class_index = classes.index(name)
                    btn1 = [x_center, y_center, ww, hh]
                    new_bndbox_list[i] = [btn1,name]

            outlabs[out_img_fn] = new_bndbox_list
            # Image.fromarray(image_auged).save(out_img_fn)
            cv2.imencode('.jpg', image_auged)[1].tofile(out_img_fn)

        if index %3==0:
            json.dump(outlabs, open(outlabs_fn, 'w'))

    json.dump(outlabs, open(outlabs_fn, 'w'))
    return dict2list(outlabs)


def img2jpg(org_dir):
    import imghdr
    imgType_list = ['jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif']
    for file in os.listdir(org_dir):
        if '.' in file:
            stem, suffix = os.path.splitext(file)
            if suffix in imgType_list:
                old_img_name = os.path.join(org_dir, file)
                new_img_name = os.path.join(org_dir, stem + ".jpg")
                os.rename(old_img_name, new_img_name)

from common.utils import *

def load_xmls(pa):
    #os.chdir(pa)
    #predict, classes = get_predict(pa, model_type)
    li = listdir(pa, ['.jpg', '.bmp', '.png', '.jpeg'])
    #os.system('del out\\*.jpg')
    # print([(872.5, 754.5, 21.0, 50.0, 266.09950256347656, '型号'), (966.0, 752.5, 24.0, 106.0, 91.11239969730377, 'TWZOOUM'), (894.0000610351562, 773.5, 23.0, 96.073, 23.0, 96.0, 267.4831237792969, '有线输入'), (964.5000610351562, 772.4999389648438, 23.0, 35.0, 91.78991055488586, '5V'), (1018.3718508911132812, 27651367188, 777.8508911132812, 28.0, 51.0, 100.61965465545654, '0.5A'), (1073.0001220703125, 786.5, 26.0, 59.0, 102.75753211975098, 'MAX'), (914.77490234375, 796.994140625, 24.0, 143.0, 268.77242279052734, '支持无线输入'), (890.5, 817.5, 31.0, 95.0, 90.0, '新定能量'), (988.76516628265381, 0.176513671875, 823.2059326171875, 31.0, 84.0, 97.76516628265381, '215wh')])
    outli = []
    for i in range(len(li)):
        fn = li[i]
        fn1, ext = os.path.splitext(fn)
        xmlfn = f'{pa}/{fn1}.xml'
        imgfn = f'{pa}/{fn}'
        if not os.path.exists(xmlfn):
            continue
        labs = read_xml_annotation(xmlfn, True)
        if len(labs)>0:
            outli.append([imgfn, labs])
    
    return outli

if __name__ == "__main__":
    org_dir = 'F:/data/210417高杰木板结巴检测/imgs/210506pic'
    org_dir = 'E:/data/210417jieba/imgs/210506pic'
    org_dir = 'D:/code/git/ywlydd/deepgui/yolov5_xx/imgs/dp_ocr'
    org_dir = 'D:/data/211105宜美哲活塞环/220114/train'
    org_dir = 'D:/data/211105宜美哲活塞环/220114/train/minout'
    org_dir = 'D:/data/201211美的纸箱喷码/train_pic_yolo'
    out_dir = f'{org_dir}/work_pa/yolov5/test'
    org_dir = 'D:/data/220329竹筷/标注caise/mini5'
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    # 将图片格式统一为了jpg
    outlabs = augmentation(org_dir, out_dir, 2, 20)
    print(len(outlabs))
